####    构建vnet


import torch
import torch.nn as nn
import vnet_def

# 动态设备分配
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class VNet(nn.Module):
    def __init__(self, num_classes=2):
        super(VNet, self).__init__()

        self.layer0 = vnet_def.res_block(1, 16, "forward0").to(device)
        self.layer11 = vnet_def.res_block(16, 32, "deconv").to(device)
        self.layer2 = vnet_def.res_block(32, 32, "forward2").to(device)
        self.layer22 = vnet_def.res_block(32, 64, "deconv").to(device)
        self.layer3 = vnet_def.res_block(64, 64, "forward3").to(device)
        self.layer33 = vnet_def.res_block(64, 128, "deconv", dropout=True).to(device)
        self.layer4 = vnet_def.res_block(128, 128, "forward3").to(device)
        self.layer44 = vnet_def.res_block(128, 256, "deconv", dropout=True).to(device)
        self.layer5 = vnet_def.res_block(256, 256, "forward3").to(device)
        self.layer55 = vnet_def.res_block(256, 256, "upconv").to(device)
        self.layer6 = vnet_def.res_block(256, 256, "forward3").to(device)
        self.layer66 = vnet_def.res_block(256, 128, "upconv").to(device)
        self.layer7 = vnet_def.res_block(128, 128, "forward3").to(device)
        self.layer77 = vnet_def.res_block(128, 64, "upconv").to(device)
        self.layer8 = vnet_def.res_block(64, 64, "forward2").to(device)
        self.layer88 = vnet_def.res_block(64, 32, "upconv").to(device)
        self.layer9 = vnet_def.res_block(32, 32, "forward1").to(device)
        self.layer10 = vnet_def.res_block(32, num_classes, "forward10").to(device)
        self.softmax = nn.Softmax(dim=1).to(device)
        self.dropv = nn.Dropout3d().to(device)

    def forward(self, x):
        out = self.layer0(x)
        link1 = out
        out = self.layer11(out)
        out = self.layer2(out)
        link2 = out
        out = self.layer22(out)
        out = self.layer3(out)
        link3 = out
        out = self.layer33(out)
        out = self.layer4(out)
        link4 = out
        out = self.layer44(out)
        out = self.layer5(out)
        out = self.layer55(out)
        out = torch.cat((self.dropv(link4), out), 1)
        out = self.layer6(out)
        out = self.layer66(out)
        out = torch.cat((self.dropv(link3), out), 1)
        out = self.layer7(out)
        out = self.layer77(out)
        out = torch.cat((self.dropv(link2), out), 1)
        out = self.layer8(out)
        out = self.layer88(out)
        out = torch.cat((self.dropv(link1), out), 1)
        out = self.layer9(out)
        out = self.layer10(out)
        out = self.softmax(out)
        return out